import asyncio
from typing import List, Dict

from llama_index.core.schema import NodeWithScore, BaseNode, NodeRelationship

from ..store.vector_store_interface import VectorStoreInterface
from .node_postprocessor_interface import NodePostprocessorInterface


class ContextualExpansionPostprocessor(NodePostprocessorInterface):
    """
    一个节点后处理器，用于在检索后扩展上下文。
    它会获取每个检索到的节点的相邻节点（前一个/后一个），
    从而为语言模型提供更完整的上下文信息。
    """

    def __init__(self, vector_store: VectorStoreInterface, window_size: int = 1):
        """
        初始化后处理器。

        Args:
            vector_store: 一个向量存储实例，必须支持通过 ID 高效检索节点。
            window_size: 要在每个检索到的节点的任一侧扩展的窗口大小。
                         例如，window_size=1 表示获取前一个和后一个节点。
        """
        self._vector_store = vector_store
        self._window_size = window_size

    async def postprocess_nodes(
        self, nodes: List[NodeWithScore], query_bundle, **kwargs
    ) -> List[NodeWithScore]:
        """
        异步地扩展节点列表，包含它们的相邻节点。
        """
        expanded_nodes: Dict[str, BaseNode] = {}
        
        # 首先将原始检索到的节点加入字典，以处理重叠并保留原始分数
        for node_with_score in nodes:
            expanded_nodes[node_with_score.node.node_id] = node_with_score.node

        # 异步地为每个节点获取其相邻节点
        fetch_tasks = []
        for node_with_score in nodes:
            fetch_tasks.append(self._fetch_neighbors(node_with_score.node, expanded_nodes))
        
        await asyncio.gather(*fetch_tasks)

        # 从字典中重建 NodeWithScore 列表，保持原始顺序和分数
        # 注意：这里我们只返回扩展后的节点列表，分数可以被认为是统一的或继承自原始节点
        final_nodes = []
        original_node_scores = {n.node.node_id: n.score for n in nodes}

        # 排序以保证一定的顺序性
        # 这里可以根据需要选择更复杂的排序逻辑，例如按原始分数或节点ID
        sorted_node_ids = sorted(expanded_nodes.keys())

        for node_id in sorted_node_ids:
            node = expanded_nodes[node_id]
            final_nodes.append(
                NodeWithScore(node=node, score=original_node_scores.get(node_id, 0.0))
            )

        return final_nodes

    async def _fetch_neighbors(self, node: BaseNode, node_dict: Dict[str, BaseNode]):
        """
        异步获取单个节点的相邻节点并将其添加到字典中。
        """
        neighbor_ids_to_fetch = []
        
        current_node = node
        # 收集前向邻居ID
        for _ in range(self._window_size):
            prev_rel = current_node.relationships.get(NodeRelationship.PREVIOUS)
            if not prev_rel or prev_rel.node_id in node_dict or prev_rel.node_id in neighbor_ids_to_fetch:
                break
            neighbor_ids_to_fetch.append(prev_rel.node_id)
            # 为了能连续向前查找，我们需要一个临时的节点对象，但这里简化为只查找直接邻居
            # LlamaIndex 的 BaseNode 没有提供从关系直接获取节点的方法，所以我们先收集ID
            # 这个逻辑在 LlamaIndex 的实现中更为复杂，此处为简化版
            break # 简化为只获取直接邻居

        current_node = node
        # 收集后向邻居ID
        for _ in range(self._window_size):
            next_rel = current_node.relationships.get(NodeRelationship.NEXT)
            if not next_rel or next_rel.node_id in node_dict or next_rel.node_id in neighbor_ids_to_fetch:
                break
            neighbor_ids_to_fetch.append(next_rel.node_id)
            break # 简化为只获取直接邻居

        if neighbor_ids_to_fetch:
            # 使用我们新增的 aget_nodes 方法批量获取所有邻居节点
            fetched_neighbors = await self._vector_store.aget_nodes(neighbor_ids_to_fetch)
            for neighbor_node in fetched_neighbors:
                node_dict[neighbor_node.node_id] = neighbor_node 